#pragma once

#include <algorithm/analytical_model/algorithm/cuda/math/vector/Common.cuh>
#include <algorithm/shared/utils/GenericVectorMath.h>
#include <utils/config.h>

#include <cuda_bf16.h>

// Not fully tested
// Not using templates for now

namespace solar::cuda
{
    using bfloat16 = __nv_bfloat16;
    using bfloat162 = __nv_bfloat162;
    struct bfloat163
    {
            bfloat16 x;
            bfloat16 y;
            bfloat16 z;
    };
} // namespace solar::cuda

inline HOT_FUNCTION __host__ __device__ auto make_bfloat163(solar::cuda::bfloat16 a, solar::cuda::bfloat16 b,
                                                            solar::cuda::bfloat16 c) -> solar::cuda::bfloat163
{
    return {a, b, c};
}

inline HOT_FUNCTION __host__ __device__ auto make_bfloat163(solar::cuda::bfloat162 a, solar::cuda::bfloat16 b) -> solar::cuda::bfloat163
{
    return {a.x, a.y, b};
}

inline HOT_FUNCTION __host__ __device__ auto make_bfloat163(solar::cuda::bfloat16 a, solar::cuda::bfloat162 b) -> solar::cuda::bfloat163
{
    return {a, b.x, b.y};
}

inline HOT_FUNCTION __host__ __device__ auto make_bfloat163(const float3& a) -> solar::cuda::bfloat163
{
    return {a.x, a.y, a.z};
}

// addition
inline HOT_FUNCTION __host__ __device__ auto operator+(const solar::cuda::bfloat163& a,
                                                       const solar::cuda::bfloat163& b) -> solar::cuda::bfloat163
{
    return {a.x + b.x, a.y + b.y, a.z + b.z};
}

inline HOT_FUNCTION __host__ __device__ void operator+=(solar::cuda::bfloat163& a, const solar::cuda::bfloat163& b)
{
    a.x += b.x;
    a.y += b.y;
    a.z += b.z;
}

// subtract
inline HOT_FUNCTION __host__ __device__ auto operator-(const solar::cuda::bfloat163& a,
                                                       const solar::cuda::bfloat163& b) -> solar::cuda::bfloat163
{
    return {a.x - b.x, a.y - b.y, a.z - b.z};
}

inline HOT_FUNCTION __host__ __device__ void operator-=(solar::cuda::bfloat163& a, const solar::cuda::bfloat163& b)
{
    a.x -= b.x;
    a.y -= b.y;
    a.z -= b.z;
}

// multiply
inline HOT_FUNCTION __host__ __device__ auto operator*(const solar::cuda::bfloat163& a,
                                                       const solar::cuda::bfloat16& s) -> solar::cuda::bfloat163
{
    return {a.x * s, a.y * s, a.z * s};
}
inline HOT_FUNCTION __host__ __device__ auto operator*(const solar::cuda::bfloat16& s,
                                                       const solar::cuda::bfloat163& a) -> solar::cuda::bfloat163
{
    return {a.x * s, a.y * s, a.z * s};
}
inline HOT_FUNCTION __host__ __device__ void operator*=(solar::cuda::bfloat163& a, const solar::cuda::bfloat16& s)
{
    a.x *= s;
    a.y *= s;
    a.z *= s;
}

// divide
inline HOT_FUNCTION __host__ __device__ auto operator/(const solar::cuda::bfloat163& a,
                                                       const solar::cuda::bfloat163& b) -> solar::cuda::bfloat163
{
    return {a.x / b.x, a.y / b.y, a.z / b.z};
}
inline HOT_FUNCTION __host__ __device__ auto operator/(const solar::cuda::bfloat163& a,
                                                       const solar::cuda::bfloat16& s) -> solar::cuda::bfloat163
{
    // const auto inv = bfloat16(1.0F) / s;
    return {a.x / s, a.y / s, a.z / s};
}
inline HOT_FUNCTION __host__ __device__ void operator/=(solar::cuda::bfloat163& a, const solar::cuda::bfloat16& s)
{
    // const auto inv = bfloat16(1.0F) / s;
    a.x /= s;
    a.y /= s;
    a.z /= s;
}

// lerp
inline HOT_FUNCTION __device__ __host__ auto lerp(const solar::cuda::bfloat163& a, const solar::cuda::bfloat163& b,
                                                  const solar::cuda::bfloat16& t) -> solar::cuda::bfloat163
{
    return a + t * (b - a);
}

// clamp
inline HOT_FUNCTION __device__ __host__ auto clamp(const solar::cuda::bfloat16& f, const solar::cuda::bfloat16& a,
                                                   const solar::cuda::bfloat16& b) -> solar::cuda::bfloat16
{
    return __hmax(a, __hmin(f, b));
}

inline HOT_FUNCTION __device__ __host__ auto clamp(const solar::cuda::bfloat163& v, const solar::cuda::bfloat16& a,
                                                   const solar::cuda::bfloat16& b) -> solar::cuda::bfloat163
{
    return {clamp(v.x, a, b), clamp(v.y, a, b), clamp(v.z, a, b)};
}

inline HOT_FUNCTION __device__ __host__ auto clamp(const solar::cuda::bfloat163& v, const solar::cuda::bfloat163& a,
                                                   const solar::cuda::bfloat163& b) -> solar::cuda::bfloat163
{
    return {clamp(v.x, a.x, b.x), clamp(v.y, a.y, b.y), clamp(v.z, a.z, b.z)};
}

// dot product
inline HOT_FUNCTION __host__ __device__ auto dot(const solar::cuda::bfloat163& a, const solar::cuda::bfloat163& b) -> solar::cuda::bfloat16
{
    return a.x * b.x + a.y * b.y + a.z * b.z;
}

// length
inline HOT_FUNCTION __host__ __device__ auto length(const solar::cuda::bfloat163& r) -> solar::cuda::bfloat16
{
    return hsqrt(dot(r, r));
}

inline HOT_FUNCTION __host__ __device__ auto rlength(const solar::cuda::bfloat163& v) -> solar::cuda::bfloat16
{
    return hrsqrt(dot(v, v));
}

// normalize
inline HOT_FUNCTION __host__ __device__ auto normalize(const solar::cuda::bfloat163& v) -> solar::cuda::bfloat163
{
    return v * rlength(v);
}

// floor
inline HOT_FUNCTION __host__ __device__ auto floor(const solar::cuda::bfloat163& v) -> solar::cuda::bfloat163
{
    return {hfloor(v.x), hfloor(v.y), hfloor(v.z)};
}

// absolute value
inline HOT_FUNCTION __host__ __device__ auto fabs(const solar::cuda::bfloat163& v) -> solar::cuda::bfloat163
{
    return {__habs(v.x), __habs(v.y), __habs(v.z)};
}

// negate
inline HOT_FUNCTION __host__ __device__ auto operator-(const solar::cuda::bfloat163& a) -> solar::cuda::bfloat163
{
    return {-a.x, -a.y, -a.z};
}

// min
static inline HOT_FUNCTION __host__ __device__ auto fminf(const solar::cuda::bfloat163& a,
                                                          const solar::cuda::bfloat163& b) -> solar::cuda::bfloat163
{
    return {__hmin(a.x, b.x), __hmin(a.y, b.y), __hmin(a.z, b.z)};
}

// max
static inline HOT_FUNCTION __host__ __device__ auto fmaxf(const solar::cuda::bfloat163& a,
                                                          const solar::cuda::bfloat163& b) -> solar::cuda::bfloat163
{
    return {__hmax(a.x, b.x), __hmax(a.y, b.y), __hmax(a.z, b.z)};
}

// cross product
inline HOT_FUNCTION __host__ __device__ auto cross(const solar::cuda::bfloat163& a,
                                                   const solar::cuda::bfloat163& b) -> solar::cuda::bfloat163
{
    return {a.y * b.z - a.z * b.y, a.z * b.x - a.x * b.z, a.x * b.y - a.y * b.x};
}

// reflect
inline HOT_FUNCTION __host__ __device__ auto reflect(const solar::cuda::bfloat163& i,
                                                     const solar::cuda::bfloat163& n) -> solar::cuda::bfloat163
{
    return i - solar::cuda::bfloat16(2.0F) * n * dot(n, i);
}
